#!/usr/bin/env python3
import argparse
import socket
import sys
import time
from scapy.all import Ether, IP, UDP, TCP, Raw, sendp
from scapy.packet import Packet, bind_layers
from scapy.fields import BitField, XShortField

# 定义与receive.py相同的SourceRoute类
class SourceRoute(Packet):
    fields_desc = [ BitField("bos", 0, 1),
                   BitField("port", 0, 15)]

class SourceRoutingTail(Packet):
    fields_desc = [ XShortField("etherType", 0x800)]

# 绑定层关系
bind_layers(Ether, SourceRoute, type=0x1234)
bind_layers(SourceRoute, SourceRoute, bos=0)
bind_layers(SourceRoute, SourceRoutingTail, bos=1)
bind_layers(SourceRoutingTail, IP)

def build_source_route_stack(ports):
    """构建源路由栈"""
    if not ports:
        return None
    
    # 创建源路由栈
    bos_list = [0] * len(ports)
    bos_list[-1] = 1  # 最后一个条目bos=1
    
    # 构建源路由头部链
    current_packet = None
    for i, (port, bos) in enumerate(zip(ports, bos_list)):
        if i == 0:
            current_packet = SourceRoute(bos=bos, port=port)
        else:
            current_packet = current_packet / SourceRoute(bos=bos, port=port)
    
    return current_packet

def send_ipv4_packet(iface, src_ip, dst_ip, src_mac, dst_mac, message="IPv4 Packet"):
    """发送普通IPv4数据包"""
    ether = Ether(src=src_mac, dst=dst_mac, type=0x800)
    ip = IP(src=src_ip, dst=dst_ip, ttl=64)
    udp = UDP(sport=1234, dport=4321)
    data = Raw(load=message.encode())
    
    packet = ether/ip/udp/data
    sendp(packet, iface=iface, verbose=False)
    print(f"Sent IPv4 packet: {src_ip} -> {dst_ip}")

def send_source_routed_packet(iface, src_ip, dst_ip, src_mac, dst_mac, route_ports, message="Source Routed Packet"):
    """发送源路由数据包（匹配receive.py格式）"""
    # 构建源路由栈
    src_route_stack = build_source_route_stack(route_ports)
    
    # 以太网帧，使用自定义类型标识源路由
    ether = Ether(src=src_mac, dst=dst_mac, type=0x1234)
    # IP包
    ip = IP(src=src_ip, dst=dst_ip, ttl=64)
    # UDP包（匹配receive.py的过滤器）
    udp = UDP(sport=1234, dport=4321)
    data = Raw(load=message.encode())
    
    # 组装数据包
    if src_route_stack:
        packet = ether/src_route_stack/ip/udp/data
    else:
        packet = ether/ip/udp/data
        
    sendp(packet, iface=iface, verbose=False)
    print(f"Sent source routed packet via ports {route_ports}")

def main():
    parser = argparse.ArgumentParser(description='Packet Sender')
    parser.add_argument('--iface', type=str, required=True, help='Interface to send packets')
    parser.add_argument('--mode', choices=['ipv4', 'source-route'], required=True, 
                       help='Packet mode: ipv4 or source-route')
    parser.add_argument('--src-ip', type=str, default='10.0.1.1', help='Source IP address')
    parser.add_argument('--dst-ip', type=str, default='10.0.2.2', help='Destination IP address')
    parser.add_argument('--src-mac', type=str, default='08:00:00:00:01:11', help='Source MAC address')
    parser.add_argument('--dst-mac', type=str, default='08:00:00:00:01:00', help='Destination MAC address')
    parser.add_argument('--route-ports', type=int, nargs='+', help='Source route ports (for source-route mode)')
    parser.add_argument('--count', type=int, default=1, help='Number of packets to send')
    parser.add_argument('--interval', type=float, default=1.0, help='Interval between packets')
    
    args = parser.parse_args()
    
    for i in range(args.count):
        if args.mode == 'ipv4':
            send_ipv4_packet(args.iface, args.src_ip, args.dst_ip, 
                           args.src_mac, args.dst_mac, 
                           f"IPv4 Packet {i+1}")
        elif args.mode == 'source-route':
            if not args.route_ports:
                print("Error: --route-ports required for source-route mode")
                sys.exit(1)
            send_source_routed_packet(args.iface, args.src_ip, args.dst_ip,
                                    args.src_mac, args.dst_mac,
                                    args.route_ports,
                                    f"Source Routed Packet {i+1}")
        
        if i < args.count - 1:
            time.sleep(args.interval)

if __name__ == '__main__':
    main()